import cv2
import os
import sys
from ultralytics import YOLO
import numpy as np
import pandas as pd
import socket
import time
from datetime import datetime

index = {'溶解氧': 'DO', 'pH值': 'pH', '温度': 'Temper', '浊度': 'Turb', '总磷': 'TP', '总氮': 'TN', '水体透明度': 'WT'}


def recvall(sock, count):
    buf = b''
    while count:
        newbuf = sock.recv(count)
        if not newbuf:
            return None
        buf += newbuf
        count -= len(newbuf)
    return buf


class Process():
    # 加载训练好的模型
    def __init__(self, model):
        self.trained_model = YOLO("../model/" + model)  # 调用模型
        self.observes = []

    def __call__(self, dur_time, list_, f):
        result = []
        index_list = list_[0]
        camera_list = ['index'] + list_[1]
        print(camera_list)

        for camera in camera_list:
            self.img = self.get_images()
            res1 = self.process_images(camera)
            res2 = self.get_other_index(index_list, camera)
            if camera == 'index':
                current_time = 'end_time'
            else:
                current_time = datetime.now()
            row = [camera] + res1 + [current_time, dur_time, f] + res2
            result.append(row)

        columns = ['Camera', 'Class_0', 'Class_1', 'Class_2', 'Class_3', 'Timestamp', 'Duration',
                   'Frame_rate'] + index_list
        result_df = pd.DataFrame(result, columns=columns)

        # 转置表格
        result_df = result_df.T
        print(result_df)

        # 保存为xlsx文件
        result_df.to_excel("result.xlsx", index=True)
        return result_df  # 返回csv格式的结果文件

    def get_images(self):
        pic_files = [f for f in os.listdir('../testpicture') if f.endswith('.jpg')]
        imgs = []
        for i in pic_files:
            file_path = '../testpicture/' + i
            imgs.append(cv2.imread(file_path))
        return imgs

    def get_other_index(self, index_list, camera):
        output = []
        if camera == 'index':
            for i in index_list:
                output.append(index[i])
            return output
        for i in index_list:
            if i == '溶解氧':
                output.append(round(np.random.uniform(10),2))
            if i == 'pH值':
                output.append(round(np.random.uniform(14),2))
            if i == '温度':
                output.append(round(np.random.uniform(50),2))
            if i == '浊度':
                output.append(round(np.random.uniform(30),2))
            if i == '总磷':
                output.append(round(np.random.uniform(0,1),2))
            if i == '总氮':
                output.append(round(np.random.uniform(0,1),2))
            if i == '水体透明度':
                output.append(round(np.random.uniform(50),2))
        return output

    def process_images(self, camera):
        if camera == 'index':
            return ['holothurian', 'echinus', 'scallop', 'starfish']
        image_files = self.img
        # 初始化当前图像的类别计数器
        class_counts = {i: 0 for i in range(4)}
        for img_file in image_files:
            # 读取测试图像
            if img_file is None:
                print(f"Warning: Unable to load image")
                continue

            # 使用模型进行推理
            results = self.trained_model(img_file)

            # 遍历检测结果，统计每种类别的数量
            for result in results:
                for box in result.boxes:
                    class_id = int(box.cls)  # 获取类别ID
                    class_counts[class_id] += 1

        # 计算每个类别的总和
        total_counts = [class_counts[i] for i in range(4)]
        return total_counts